Repository navigation
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces several performance optimizations for Wan models on TPU, including a fused short-KV cross-attention Pallas kernel, token-sharded patch embedding, inverse Ulysses all-to-all optimizations, and token sequence padding. A critical issue was identified in the new Pallas cross-attention kernel where the BlockSpec index mapping fails to scale grid indices by their block sizes, which would result in incorrect memory slicing and wrong outputs.
2ba8966 to
e86f174
Compare
e86f174 to
1d4c9f6
Compare
0d91078 to
4703103
Compare
295c3b5 to
7a60c3d
Compare
08d405c to
7db8313
Compare
7db8313 to
6fdc94e
Compare
ebaac14 to
58f889e
Compare
58f889e to
68a580f
Compare
68a580f to
395f932
Compare
cdc584f to
2eba645
Compare
08ccfb2 to
68c3732
Compare
|
Thanks @eltsai — ran the sweep at 41 / 81 / 121 / 161 frames (2.5 / 5 / 7.5 / 10 s @ 16 fps), 720p, 40 steps, seed 12345, same prompt, DP=2×CP=4, DVFS unpinned, on tpu7x-8 and v6e-8. A =
No regression vs Note that all runs use the launcher's 81 f-tuned flash block sizes (v6e: BQ 9472 / BKV 1024; tpu7x: BQ 4736 / BKV 2048). We did not do a block-size sweep per length: the sequence length (and therefore the per-shard length and its alignment to BQ) differs for each frame count, so the optimal block sizes need to be found separately for each frame choice. The non-81 f numbers above are therefore conservative for all three configurations. * v6e @ 121 f is the one length where the launcher's v6e profile (Ulysses U=4, BQ 9472) is the wrong pick: the per-shard length (27,900 tokens) is not block-aligned, and the 2D-ring recipe that is also in this stack ( All runs were bit-reproducible across cold/warm AOT (identical md5s), and the outputs are visually clean on both platforms. |
0414334 to
34789ff
Compare
34789ff to
87c36b2
Compare
mbohlool
left a comment
There was a problem hiding this comment.
a requested lane mode should raise rather than turn itself off, and CI should run at least one masking test.
Last of four pieces split out of the Wan fast-inference optimizations PR
(the cross-attention Pallas kernel, token-sharded patch embedding and
shard-major Ulysses A2A are the three PRs below this one). The switch
defaults off in the Wan YAMLs and in wan_runtime_options; it is a plain
config key plumbed to WanModel through `attention_config_entries(config)`.
- Lane sequence padding (`wan_seq_pad="lane"`): `WanModel` pads the video tokens
at the tail to a per-shard multiple of 128 (`75,600 -> 75,776` at
720p/CP=4; on the Ulysses-ring path each ring segment is padded at its own
tail) so the Ulysses reshapes are copy-free, pads the rotary embedding to
match, and unpads the output and residual. The padding descriptor
(`TokenPadding`) is planned via `plan_token_padding` and passed explicitly as
a `token_padding` argument from `WanModel.__call__` ->
`WanTransformerBlock.__call__` -> `FlaxWanAttention.__call__` ->
`NNXAttentionOp.apply_attention` (no process/thread-global `ContextVar`).
- The custom `ulysses_custom*` / `ulysses_ring_custom*` kernels mask the pad
keys by passing the real length as the ragged KV tail; the fixed-m norm
reductions (`_compute_fixed_m_metadata` via `q_valid_len` on both the 1-D
Ulysses path and the R == 1 ring path, `_ring_fixed_m_norms_pre_a2a` via
`row_ok` for R > 1) exclude the pad Q/K/V rows so pad-token values cannot
push a block onto the online-softmax path or move the `k_mean`. The ring's
`ring_kv_seq_len` is always the real length.
- Fails closed on invalid static/mode configs: when `wan_seq_pad="lane"` is
requested, incompatible static configurations (`per_token_t=True` TI2V
timesteps, `use_svg_attention`, `use_memory_efficient_attention`, unsupported
kernels, or `ulysses_shards` that do not divide the context axis) raise
`ValueError` in `plan_token_padding` (checked at `WanModel.__init__` time as
well as per call); `_resolve_token_padding` raises `ValueError` on any
`TokenPadding.total_len != seq_len` mismatch; and attention kernels that
cannot mask active token padding raise `NotImplementedError` via
`_reject_token_padding`. Dynamic short (`S < flash_min_seq_length`) or
indivisible (`S % context_shards != 0`) sequences log a warn-once message and
run unpadded.
- `wan_block_benchmark.py` applies `plan_token_padding` to the benchmarked
sequence length and passes `token_padding` explicitly to the block.
- `run_wan_fast_inference.sh`: the v6e and v7 profiles set `WAN_SEQ_PAD=lane`
and `DEFAULT_ULYSSES_OUT_A2A=shard_major` (overridable); the generic profile
keeps `off` and `flat`. LIBTPU: v6e disables
`--xla_tpu_enable_ici_ag_pipelining` / `--xla_tpu_enable_ag_backward_pipelining`,
v7 enables them and adds `--xla_tpu_enable_offloading_copy_to_sparsecore=false`
and `BQ=4736`.
Performance, measured at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps,
warm AOT, DVFS unpinned unless noted; CP=4, DP=2, launcher defaults):
v6e-8: denoise 124.5s - 125.0s, generate 126.6s - 127.5s
tpu7x-8: denoise 104.9s, generate 107.8s (106.6s / 109.6s with fused RoPE off)
tpu7x-8, DVFS pinned (`PIN_DVFS_P_STATE=true`):
denoise 91.7s, generate 94.6s with the fused RoPE kernel on
denoise 93.4s, generate 96.4s with the fused RoPE kernel off
Tests
- `wan_seq_pad_test.py` (19): padding plan, `ValueError` guards for invalid
static/mode configs (including `per_token_t=True`) and `_resolve_token_padding`
length mismatches, warn-once logging for short/indivisible sequences, pad-key
masking for `ulysses_custom*` and `ulysses_ring_custom*` (U=2 x R=2, U=4 x R=1,
and shard-major U=2 x R=2), fixed-m metadata ignoring pad query rows,
`NotImplementedError` kernel guards, and end-to-end `WanModel` lane padding
parity. All 19 tests run unconditionally in GitHub Actions (no
`_SKIP_IN_GITHUB_ACTIONS`).
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
validation; `wan_seq_pad` is excluded from the `FlaxWanAttention` attribute
check since `WanModel` consumes it.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `wan_seq_pad_test.py`.
Stacked on
feat/ulysses-shard-major-a2a. Last of four pieces split out of the original #491 (this PR keeps the number). Addswan_seq_pad: "lane"(default"off"):WanModelpads the video tokens at the tail to a per-shard multiple of 128 (75,600 → 75,776at 720p/CP=4; per ring segment on the Ulysses-ring path) so the Ulysses reshapes are copy-free, pads the rotary embedding to match, and unpads the output/residual. The padding is declared to the attention layers withself_attention_token_padding(TokenPadding, aContextVarscoped to the block trace); the customulysses/ulysses_ring_customkernels mask the pad keys via the ragged KV tail, and the fixed-m norm reductions (q_valid_lenon both the 1-D Ulysses path and theR == 1ring path,row_okforR > 1) exclude pad Q/K/V rows so they cannot push a block onto the online-softmax path or movek_mean. Every configuration that disables padding logs why, once per cause. Behaviour change: with padding active, paths that cannot mask it (flash/ non-customulysses,ulysses_ring,dot_product,cudnn_flash_te,head_local_svg,ring bidir) raiseNotImplementedErrorinstead of silently attending to pads.run_wan_fast_inference.shsetslanefor v6e/v7 and carries the LIBTPU flag changes (ici_ag/ag_backwardpipelining off on v6e, on on v7; v7 adds--xla_tpu_enable_offloading_copy_to_sparsecore=falseandBQ=4736).Performance (Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2)Per-PR progression across the 10-PR stack on TPU v6e-8 (DVFS unpinned)
main(f1d0fafb)origin/main(1D UlyssesU=4/ 2D RingU=2, R=2)feat/fixed-m-kernel(per_q_block,U=4)mainfeat/attn-dot-fallback-fixesfeat/ring-attention(U=4/ 2D RingU=2, R=2)mainfeat/fused-rmsnorm-rope-xlafeat/wan-fast-serving(cold / warm AOT)feat/wan-custom-kernels(Pallas fused RoPE ON, cold / warm)feat/wan-cross-attn-pallas(cold / warm AOT)feat/wan-token-patch-embed(cold / warm AOT)feat/ulysses-shard-major-a2a(wan_ulysses_out_a2a="chunked")feat/wan-fast-inference-optimizations(wan_seq_pad="lane",WAN_ULYSSES_OUT_A2A=flatablation)feat/wan-fast-inference-optimizations(full stack, cold / warm AOT)main)USE_FUSED_ROPE_KERNEL=false.Tests:
wan_seq_pad_test.py(18, new; padding plan + disable logging, pad-key masking acrossulyssescustom /ulysses_ring_customincl.U=2×R=2,U=4×R=1, and the scaled-down tpu7xU=2×R=2shard-major path, fixed-m metadata ignoring pad rows,NotImplementedErrorguards),wan_runtime_options_test.py(+enum validation).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491